#Nature's Kidneys: the Role of Wetland Reserve Easements in Restoring Water Quality
#Nicole Karwowski and Marin Skidmore
#Replication File
#2025

#Main results

#Packages
library(lubridate)
library(ggplot2)
library(dplyr)
library(fixest)
library(plm)
library(ggplot2)
library(readr)
library(mapview)
library(fixest)
library(stringr)
library(stargazer)
library(cowplot)
library(broom)
library(openxlsx)
library(boot)
library(pvaluefunctions)
library(ggfixest)
library(patchwork)

#set wd
setwd("C:/Users/n57m645/OneDrive - Montana State University/UW_Madison/Water quality/Replication Package")

#load main dataset
load("Data/main/main_huc12_panel.RData")

#keep only relevant years 
merge<-subset(merge, as.numeric(Year)>=1985)

#keep only observations that have an ammonia or phosphorous measurement
merge<-subset(merge, !is.na(ammonia_filtered) | !is.na(TP_filtered) | !is.na(TKN_filtered))

#give ever-treated their own df 
treated <- subset(merge, merge$EverTreated==1)
set.seed(10)

#################################################################################


#Create variable labels for tables
fixest::setFixest_dict(c(ammonia_filtered="Ammonia",
                         TP_filtered="Phosphorus",
                         TKN_filtered="Total Kjedahl Nitrogen",
                         RestoredBinary="Restored wetland (0/1)",
                         RestoredCumP="Restored wetland (prop)",
                         RestoredCum="Restored wetland (acres)",
                         TP_filtered_ihs="IHS(phosphorus)",
                         ammonia_filtered_ihs="IHS(ammonia)",
                         TKN_filtered_ihs="IHS(TKN)",
                         ihs_restored="IHS restored wetland (acres)",
                         ihs_upstream_restored_sf="IHS upstream restored wetland (acres)",
                         RestoredCumP_ihs="IHS restored wetland (prop)",
                         RestoredCum_ihs="IHS restored wetland (acres)",
                         UpstreamRestoredCumSfP="Upstream restored wetland (prop)",
                         ihs_upstream_restored_sf="IHS upstream restored wetland (acres)",
                         UpstreamRestoredCumSf="Upstream restored wetland (acres)",
                         RestoredBinaryUpstream="Upstream restored wetland (01/)",
                         UpstreamAmmoniaFiltered="Upstream ammonia",
                         UpstreamTpFiltered="Upstream phosphorus",
                         UpstreamTknFiltered="Upstream TKN",
                         upstream_sf_ammonia_filtered="Upstream ammonia (w=streamflow)",
                         upstream_sf_TP_filtered="Upstream phosphorus (w=streamflow)",
                         upstream_sf_TKN_filtered="Upstream TKN (w=streamflow)",
                         Ppt="Ppt",
                         Temp="Temp",
                         barren="Barren",
                         cropland="Cropland",
                         developed="Developed",
                         vegetation="Vegetation",
                         water="Water",
                         grass_shrub="Grass shrub",
                         ice_snow="Ice snow",
                         tree_cover="Tree cover",
                         AnimalUnits="Animal units",
                         tri_ammonia="Ammonia dischage",
                         tri_ammonia_water="Ammonia discharge (water)",
                         tri_nitrate="Nitrate discharge",
                         tri_nitrate_water="Nitrate discharge (water)", 
                         tri_phosphorus="Phosphorus discharge",
                         tri_phosphorus_water="Phosphorus discharge (water)",
                         huc8="HUC8",
                         huc4="HUC4",
                         huc12="HUC12",
                         huc4year="HUC4 x year",
                         huc4month="HUC4 x month",
                         Year="Year",
                         Month="Month",
                         TTTyear2="Period",
                         TTT.closedyear2 = "Period",
                         EverTreated="",
                         times=""))

# Start regressions 
#fixed effect model set up 
setFixest_fml(..ctrl_n = ~ +Ppt+Temp+cropland+developed+              
                vegetation+water + tri_ammonia_water,
              ..ctrl_p = ~ +Ppt+Temp+cropland+developed+              
                vegetation+water + tri_phosphorus_water,
              ..huc8_fe =~huc8 + Year + Month,
              ..huc12_fe = ~huc12 + Year + Month,
              ..huc4year_fe = ~huc8 + huc4year + Month,
              ..int = ~+ihs_restored*Ppt+ihs_restored*Temp+ihs_restored*cropland+ihs_restored*developed+              
                ihs_restored*vegetation+ihs_restored*water)

#post-processing functions
remove_model_numbers <- function(x) {
  # Remove lines starting specifically with & (1) & (2)
  x <- gsub("^\\s*& \\(1\\) & \\(2\\).*\n", "", x, perl = TRUE)
  return(x)
}

change_column_widths<- function(x) {
  #adjust the column widths of the table
  x <- gsub("{lcccccc}", "{m{6cm} m{1.4cm} m{1.4cm} m{1.4cm} m{1.4cm} m{1.4cm} m{1.4cm}}", x, perl=TRUE)
  return(x)
}

#function to remove model numbers and adjust column width
both<-function(x){
  x <- gsub("^\\s*& \\(1\\) & \\(2\\).*\n", "", x, perl = TRUE)
  x <- gsub("{lcccccc}", "{m{6cm} m{1.4cm} m{1.4cm} m{1.4cm} m{1.4cm} m{1.4cm} m{1.4cm}}", x, perl=TRUE)
  return(x)
}

#function to count the number of huc12 in our models
count_huc12 <- function(model, data, unit_var = "huc12") {
  # Extract observation row indices from the model
  obs_model <- obs(model)  
  
  # Ensure the dataset has a row index
  if (!"row_id" %in% names(data)) {
    data$row_id <- 1:nrow(data)  
  }
  
  # Filter dataset to match rows used in the model
  filtered_data <- data[data$row_id %in% obs_model, ]
  
  # Count unique units in the specified column
  return(length(unique(filtered_data[[unit_var]])))
}


#main table. panel with fixed effects specifications.
#effect of current wrp area (IHS acres) on huc12 water quality that ever had a wrp restoration
#panel A: ammonia
#main specification: ihs restored area and level WQ
a1<-feols(fml=ammonia_filtered~ihs_restored+..ctrl_n | ..huc8_fe, 
          cluster="huc8year", data=merge, subset=merge$EverTreated==1)

a2<-feols(fml=ammonia_filtered~ihs_restored+..ctrl_n| ..huc4year_fe,
          cluster="huc8year", data=merge, subset=merge$EverTreated==1)

a3<-feols(fml=ammonia_filtered~ihs_restored+..ctrl_n| ..huc12_fe,
          cluster="huc8year", data=merge, subset=merge$EverTreated==1)

#add upstream restored wetlands 
a1u<-feols(fml=ammonia_filtered~ihs_restored+ihs_upstream_restored_sf+..ctrl_n| ..huc8_fe,
           cluster="huc8year", data=merge, subset=merge$EverTreated==1)

a2u<-feols(fml=ammonia_filtered~ihs_restored+ihs_upstream_restored_sf+..ctrl_n| ..huc4year_fe,
           cluster="huc8year", data=merge, subset=merge$EverTreated==1)

a3u<-feols(fml=ammonia_filtered~ihs_restored+ihs_upstream_restored_sf+..ctrl_n| ..huc12_fe,
           cluster="huc8year", data=merge, subset=merge$EverTreated==1)

#Phosphorus
p1<-feols(fml=TP_filtered~ihs_restored+..ctrl_p | ..huc8_fe, 
          cluster="huc8year", data=merge, subset=merge$EverTreated==1)

p2<-feols(fml=TP_filtered~ihs_restored+..ctrl_p| ..huc4year_fe,
          cluster="huc8year", data=merge, subset=merge$EverTreated==1)

p3<-feols(fml=TP_filtered~ihs_restored+..ctrl_p| ..huc12_fe,
          cluster="huc8year", data=merge, subset=merge$EverTreated==1)

#add upstream restored wetlands 
p1u<-feols(fml=TP_filtered~ihs_restored+ihs_upstream_restored_sf+..ctrl_p| ..huc8_fe,
           cluster="huc8year", data=merge, subset=merge$EverTreated==1)

p2u<-feols(fml=TP_filtered~ihs_restored+ihs_upstream_restored_sf+..ctrl_p| ..huc4year_fe,
           cluster="huc8year", data=merge, subset=merge$EverTreated==1)

p3u<-feols(fml=TP_filtered~ihs_restored+ihs_upstream_restored_sf+..ctrl_p| ..huc12_fe,
           cluster="huc8year", data=merge, subset=merge$EverTreated==1)

#TKN
t1<-feols(fml=TKN_filtered~ihs_restored+..ctrl_n | ..huc8_fe, 
          cluster="huc8year", data=merge, subset=merge$EverTreated==1)

t2<-feols(fml=TKN_filtered~ihs_restored+..ctrl_n| ..huc4year_fe,
          cluster="huc8year", data=merge, subset=merge$EverTreated==1)

t3<-feols(fml=TKN_filtered~ihs_restored+..ctrl_n| ..huc12_fe,
          cluster="huc8year", data=merge, subset=merge$EverTreated==1)

#add upstream restored wetlands 
t1u<-feols(fml=TKN_filtered~ihs_restored+ihs_upstream_restored_sf+..ctrl_n| ..huc8_fe,
           cluster="huc8year", data=merge, subset=merge$EverTreated==1)

t2u<-feols(fml=TKN_filtered~ihs_restored+ihs_upstream_restored_sf+..ctrl_n| ..huc4year_fe,
           cluster="huc8year", data=merge, subset=merge$EverTreated==1)

t3u<-feols(fml=TKN_filtered~ihs_restored+ihs_upstream_restored_sf+..ctrl_n| ..huc12_fe,
           cluster="huc8year", data=merge, subset=merge$EverTreated==1)

################################################################################
#fixed effects table
fixest::etable(a1, a2, a3, a1u, a2u, a3u, 
               depvar = FALSE,drop.section = "fixef",
               tex=TRUE, digits=3, fitstat=c("my", "n", "ar2"), 
               keep=c("IHS"),
               replace = TRUE,
               file="Results/ammonia_treatedhuc12_FE.tex",
               postprocess.tex = change_column_widths,
               style.tex=fixest::style.tex("aer", line.bottom = "",
                                           var.title = "$\\textbf{Panel A: ammonia}$"),
               extralines=list("_^# HUC12" = sapply(list(a1, a2, a3, a1u, a2u, a3u), count_huc12, data = merge),
               "__ "=c(" ", " ", " ", " ", " ", " "),
               "^^  "=c(" ", " ", " ", " ", " ", " ")))

fixest::etable(t1, t2, t3, t1u, t2u, t3u, 
               depvar = FALSE, drop.section = "fixef",
               tex=TRUE, digits=3, fitstat=c("my", "n", "ar2"), 
               keep=c("IHS"),
               replace = TRUE,
               file="Results/tkn_treatedhuc12_FE.tex",
               postprocess.tex = both,
               style.tex=fixest::style.tex("aer", line.bottom = "\\hline",
                                           var.title = "$\\textbf{Panel B: total kjedahl nitrogen}$"),
               extralines=list("_^# HUC12" = sapply(list(t1, t2, t3, t1u, t2u, t3u), count_huc12, data = merge),
               "__ "=c(" ", " ", " ", " ", " ", " "),
               "^^  "=c(" ", " ", " ", " ", " ", " ")))


fixest::etable(p1, p2, p3, p1u, p2u, p3u, 
               depvar = FALSE, drop.section = "fixef",
               tex=TRUE, digits=3, fitstat=c("my", "n", "ar2"), 
               keep=c("*IHS"),
               replace = TRUE,
               file="Results/phosphorus_treatedhuc12_FE.tex",
               postprocess.tex = both,
               style.tex=fixest::style.tex("aer", line.top ="", line.bottom = "\\hline",
                                           var.title = "$\\textbf{Panel C: phosphorus }$"),
               extralines=list("_^# HUC12" = sapply(list(p1, p2, p3, p1u, p2u, p3u), count_huc12, data = merge),
                               "__"=c(" ", " ", " ", " "," ", " "),
                               "__Ever treated"=c("$\\checkmark$", "$\\checkmark$", "$\\checkmark$", "$\\checkmark$","$\\checkmark$", "$\\checkmark$"),
                               "__Controls"=c("$\\checkmark$", "$\\checkmark$", "$\\checkmark$", "$\\checkmark$","$\\checkmark$", "$\\checkmark$"),
                               "__"=c(" ", " ", " ", " "," ", " "),
                               "__HUC8 fixed effects"=c("$\\checkmark$",  "$\\checkmark$"," ", "$\\checkmark$", "$\\checkmark$"," "),
                               "__HUC12 fixed effects"=c(" "," ", "$\\checkmark$", " ", " ", "$\\checkmark$"),
                               "__Year fixed effects"=c("$\\checkmark$"," ", "$\\checkmark$", "$\\checkmark$"," ","$\\checkmark$"),
                               "__HUC4 x year fixed effects"=c(" ", "$\\checkmark$", " "," ", "$\\checkmark$"," "),
                               "__Month fixed effects"=c("$\\checkmark$", "$\\checkmark$", "$\\checkmark$", "$\\checkmark$","$\\checkmark$", "$\\checkmark$")))


###############################################################################
#Main event study specification

#ATT 
#with controls 
a2es_att<-fixest::feols(ammonia_filtered~sunab(first.treat.year, TTTyear2, att = TRUE, ref.c = 2018) +..ctrl_n | ..huc8_fe, 
                        data=merge, cluster="huc8year", subset=merge$EverTreated==1& merge$Year < 2018)
#without controls 
#ever treated
a2es_uc_att<-fixest::feols(ammonia_filtered~sunab(first.treat.year, TTTyear2, att = TRUE, ref.c = 2018)  | ..huc8_fe, 
                           data=merge, cluster="huc8year", subset=merge$EverTreated==1& merge$Year < 2018)

#ATT
#with controls 
n2es_att<-fixest::feols(TKN_filtered~sunab(first.treat.year, TTTyear2, att = TRUE, ref.c = 2018)+..ctrl_n  | ..huc8_fe, 
                        data=merge, cluster="huc8year", subset=merge$EverTreated==1& merge$Year < 2018)

#without controls
n2es_uc_att<-fixest::feols(TKN_filtered~sunab(first.treat.year, TTTyear2, att = TRUE, ref.c = 2018)  | ..huc8_fe, 
                           data=merge, cluster="huc8year", subset=merge$EverTreated==1 & merge$Year < 2018)

#ATT
#with controls 
p2es_att<-fixest::feols(TP_filtered~sunab(first.treat.year, TTTyear2, att = TRUE, ref.c = 2018)+..ctrl_p  | ..huc8_fe, 
                        data=merge, cluster="huc8year", subset=merge$EverTreated==1& merge$Year < 2018)

#without controls
p2es_uc_att<-fixest::feols(TP_filtered~sunab(first.treat.year, TTTyear2, att = TRUE, ref.c = 2018)  | ..huc8_fe, 
                           data=merge, cluster="huc8year", subset=merge$EverTreated==1& merge$Year < 2018)


#sun and abraham att table 
fixest::etable(a2es_att, a2es_uc_att,  n2es_att, n2es_uc_att, p2es_att, p2es_uc_att,  
                keep="ATT", interaction.combine="",
                tex=TRUE, digits=3, fitstat=c("my", "n", "ar2"),
                style.tex=fixest::style.tex("aer"),
                file="Results/both_es_att.tex",
                replace = TRUE,
               drop.section = "fixef",
                extralines=list("_^# HUC12" = sapply(list(a2es_att, a2es_uc_att,  n2es_att, n2es_uc_att, p2es_att, p2es_uc_att), count_huc12, data = merge)))

######################################################################################################
#Main event study coeffiecient plots

#ammonia all huc
aes1<-fixest::feols(ammonia_filtered~sunab(first.treat.year, TTTyear2, ref.c=2018) + ..ctrl_n | ..huc8_fe, 
                    data=merge, cluster="huc8year")


#ammonia ever treated
aes2<-fixest::feols(ammonia_filtered~sunab(first.treat.year, TTTyear2, ref.c=2018) + ..ctrl_n | ..huc8_fe, 
                    data=merge, cluster="huc8year", subset=merge$EverTreated==1)


#tkn all huc
tes1<-fixest::feols(TKN_filtered~sunab(first.treat.year, TTTyear2, ref.c=2018) + ..ctrl_n | ..huc8_fe, 
                    data=merge, cluster="huc8year")

#tkn ever treated
tes2<-fixest::feols(TKN_filtered~sunab(first.treat.year, TTTyear2, ref.c=2018) + ..ctrl_n | ..huc8_fe, 
                    data=merge, cluster="huc8year", subset=merge$EverTreated==1)

#phosphorus all huc
pes1<-fixest::feols(TP_filtered~sunab(first.treat.year, TTTyear2, ref.c=2018) + ..ctrl_p| ..huc8_fe, 
                    data=merge, cluster="huc8year")

#phosphorus ever treated
pes2<-fixest::feols(TP_filtered~sunab(first.treat.year, TTTyear2, ref.c=2018) + ..ctrl_p| ..huc8_fe, 
                    data=merge, cluster="huc8year", subset=merge$EverTreated==1)
summary(pes2, agg="cohort")


#use ggplot to make the event study graphs
My_Theme = theme(title=element_text(size=14),
                 text = element_text(color = "black"),
                 strip.text = element_text(color = "black", size=14),
                 strip.background=element_rect(colour=NA,
                                               fill=NA),
                 axis.title.x = element_text(size = 14),
                 axis.text.x = element_text(size = 14),
                 axis.text.y = element_text(size = 14),
                 axis.title.y = element_text(size = 14),
                 legend.position = "none")

#variable names
var_labels<-as_labeller(c('ammonia_filtered'='Ammonia', 'TP_filtered'='Phosphorus', 'TKN_filtered'='Total Kjeldahl nitrogen',
                          'All HUC12'= 'All HUC12', 'Ever-treated HUC12'= 'Ever-treated HUC12'))

#plot
es_ammonia<-ggiplot(list("All HUC12"=aes1, "Ever-treated HUC12"=aes2), 
                    multi_style="facet",
                    facet_args=list(ncol=2, labeller=var_labels),
                    main="Ammonia",
                    ylab="Estimate and 95% CI",
                    pt.join=TRUE,
                    geom_style = 'ribbon',
                    pt.col=c("black", "black","black", "black"), col=c("black","black","black", "black"),
                    xlab="",)+My_Theme+scale_y_continuous(breaks=seq(-0.2, 0.2, by=0.1))+coord_cartesian(ylim=c(-.2, 0.2))
es_ammonia

es_tkn<-ggiplot(list("All HUC12"=tes1, "Ever-treated HUC12"=tes2), 
                multi_style="facet",
                facet_args=list(ncol=2, labeller=var_labels),
                main="Total kjedahl nitrogen",
                ylab="Estimate and 95% CI",
                pt.join=TRUE,
                geom_style = 'ribbon',
                pt.col=c("black", "black","black", "black"), col=c("black","black","black", "black"),
                xlab="",)+My_Theme+scale_y_continuous(breaks=seq(-0.4, 0.4, by=0.2))+coord_cartesian(ylim=c(-.4, 0.4))
es_tkn

es_phosphorus<-ggiplot(list("All HUC12"=pes1, "Ever-treated HUC12"=pes2), 
                       multi_style="facet",
                       facet_args=list(ncol=2, labeller=var_labels),
                       main="Phosphorus",
                       ylab="Estimate and 95% CI",
                       pt.join=TRUE,
                       geom_style = 'ribbon',
                       pt.col=c("black", "black","black", "black"), col=c("black","black","black", "black"),
                       xlab="Time to treatment (year)",)+My_Theme+scale_y_continuous(breaks=seq(-0.2, 0.2, by=0.1))+coord_cartesian(ylim=c(-.2, 0.2))
es_phosphorus

#stack them together and save
es_all<-es_ammonia/es_tkn/es_phosphorus
es_all

ggsave(es_all, file="Figures/es_main_controls_sa.eps",  width=9, height=12, units="in", device = grDevices::cairo_ps)


####################################################################################
#Figure 3. Heterogeneity analysis by cropland quartiles

###################################################################################
#Heterogeneity analysis 

#ag categories 
ever_treated <- subset(merge, EverTreated == 1)
quantile(ever_treated$cropland, probs = c(.33, .67))
ever_treated$cropland_quartile <- cut(ever_treated$cropland, breaks = 3)
ag_q1 <- subset(ever_treated, cropland <= 0.25)
ag_q2 <- subset(ever_treated, cropland > 0.25 & cropland <= 0.50 )
ag_q3 <- subset(ever_treated, cropland > 0.50 & cropland <=0.75 )
ag_q4 <- subset(ever_treated, cropland > 0.75 )

tkn_q2 <- subset(ag_q2, !is.na(TKN_filtered))
tkn_q4 <- subset(ag_q4, !is.na(TKN_filtered))
length(unique(tkn_q2$huc12))
length(unique(tkn_q4$huc12))


#manually add the means in 
summary(ag_q1$ammonia_filtered)
summary(ag_q2$ammonia_filtered)
summary(ag_q3$ammonia_filtered)
summary(ag_q4$ammonia_filtered)

summary(ag_q1$TP_filtered)
summary(ag_q2$TP_filtered)
summary(ag_q3$TP_filtered)
summary(ag_q4$TP_filtered)

summary(ag_q1$TKN_filtered)
summary(ag_q2$TKN_filtered)
summary(ag_q3$TKN_filtered)
summary(ag_q4$TKN_filtered)

#this was copied into the latex table after printing 
#Dependent variable mean (ammonia) & 0.065 & 0.065 & .065 & 0.096 & 0.096 & 0.096 & 0.122 & 0.122 & 0.122 & 0.250 & 0.250 & 0.250 \\ 
#Dependent variable mean (phosphorus) & 0.094 & 0.094 & 0.094 & 0.095 & 0.095& 0.095  & 0.153  & 0.153 & 0.153 & 0.232   & 0.232  & 0.232 \\ 
#Dependent variable mean (TKN) & 0.310 & 0.310 & 0.310 & 0.370 & 0.370 & 0.370 & 0.571 & 0.571 & 0.571 & 0.710 & 0.710 & 0.710 \\ 
#run models on the four quartiles 
for (q in 1:4){
  data_q <- get(paste0("ag_q", q))
  model <- feols(fml=ammonia_filtered~ihs_restored +..ctrl_n| ..huc8_fe,
                 cluster="huc8year", data=data_q)
  assign(paste0("a1_q", q), model)
  model <- feols(fml=ammonia_filtered~ihs_restored+..ctrl_n| ..huc4year_fe,
                 cluster="huc8year", data=data_q)
  assign(paste0("a2_q", q), model)
  model <- feols(fml=ammonia_filtered~ihs_restored+..ctrl_n| ..huc12_fe,
                 cluster="huc8year", data=data_q)
  assign(paste0("a3_q", q), model)
  model <- feols(fml=TP_filtered~ihs_restored +..ctrl_p| ..huc8_fe,
                 cluster="huc8year", data=data_q)
  assign(paste0("p1_q", q), model)
  model <- feols(fml=TP_filtered~ihs_restored+..ctrl_p| ..huc4year_fe,
                 cluster="huc8year", data=data_q)
  assign(paste0("p2_q", q), model)
  model <- feols(fml=TP_filtered~ihs_restored+..ctrl_p| ..huc12_fe,
                 cluster="huc8year", data=data_q)
  assign(paste0("p3_q", q), model)
  model <- feols(fml=TKN_filtered~ihs_restored +..ctrl_n| ..huc8_fe,
                 cluster="huc8year", data=data_q)
  assign(paste0("n1_q", q), model)
  model <- feols(fml=TKN_filtered~ihs_restored+..ctrl_n| ..huc4year_fe,
                 cluster="huc8year", data=data_q)
  assign(paste0("n2_q", q), model)
  model <- feols(fml=TKN_filtered~ihs_restored+..ctrl_n| ..huc12_fe,
                 cluster="huc8year", data=data_q)
  assign(paste0("n3_q", q), model)
  
}

models_a <- list(a1_q1, a2_q1, a3_q1, a1_q2, a2_q2, a3_q2, a1_q3, a2_q3, a3_q3, a1_q4, a2_q4, a3_q4)
models_p <- list(p1_q1, p2_q1, p3_q1, p1_q2, p2_q2, p3_q2, p1_q3, p2_q3, p3_q3, p1_q4, p2_q4, p3_q4)
models_n <- list(n1_q1, n2_q1, n3_q1, n1_q2, n2_q2, n3_q2, n1_q3, n2_q3, n3_q3, n1_q4, n2_q4, n3_q4)

fitstats <- lapply(models_a, fitstat, "my")
#create group labels 
get_group_label <- function(index) {
  labels <- c("Low", "Medium", "Medium High", "Very High")
  return(labels[ceiling(index / 3)])
}

# Define the function
plot_coefficients <- function(models) {
  # Extract coefficients for "ihs_restored" and create the necessary data frame
  
  # Modify the code to use the custom function for group labels
  coef_data <- bind_rows(lapply(seq_along(models), function(i) {
    tidy(models[[i]], conf.int = TRUE) %>%
      filter(term == "ihs_restored") %>%
      mutate(
        Group = get_group_label(i),  # Assign custom group labels
        Model = rep(c("HUC8+year", "HUC8+HUC4xYear", "HUC12+Year"), times = 4)[i]  # Model labels
      )
  }))    
  # Define position dodge for offsets
  dodge <- position_dodge(width = 0.5)
  
  # Create and return the plot
  
  ggplot(coef_data, aes(x = Group, y = estimate, shape = Model)) +
    geom_point(aes(group = interaction(Group, Model)), position = dodge, size = 3) +
    geom_errorbar(aes(ymin = conf.low, ymax = conf.high, 
                      group = interaction(Group, Model),
                      linetype = Group), 
                  width = 0.2, position = dodge) +
    geom_hline(yintercept = 0, color = "gray") +
    theme_minimal() +
    labs(title = "", 
         x = "", y = "mg/L", 
         linetype = "Quartile Group", shape = "Model") +
    theme(axis.text.x = element_blank(),
          axis.ticks.x = element_blank(),
          axis.title.x = element_blank(),
          axis.title.y = element_text(size = 14),
          axis.text.y = element_text(size = 14))
  
}


get_group_label <- function(index) {
  labels <- c("Low", "Medium", "Medium High", "Very High")
  return(labels[ceiling(index / 3)])
}


# Example call to the function
plot <- plot_coefficients(models_a)
legend <- get_legend(plot)
legend_plot <- plot_grid(legend)
plot_a <- plot + theme(legend.position = "none")
plot <- plot_coefficients(models_n)
plot_n <- plot + theme(legend.position = "none")
plot <- plot_coefficients(models_p)
plot_p <- plot + theme(legend.position = "none")


combined_plot <- plot_grid(
  plot_a, plot_n,
  plot_p, legend_plot,
  labels = c("Ammonia", "TKN", "Phosphorus", ""),  
  label_fontface = "plain",       # Removes bold
  label_x = 0.5,                  
  label_size = 14,                
  ncol = 2, align = "hv"
)

combined_plot 

ggsave("Figures/ag_combined_plot.eps", plot = combined_plot, device = "eps", width = 8, height = 6, dpi = 300)

